import pickle, cv2, math
import numpy as np
from PIL import Image, ImageFont, ImageDraw

with open("projects/ocr/weights/temp.pkl","rb") as f:
    ocr_result = pickle.load(f)[0]

img_data = cv2.imread("data/input/ocr/zh_train_5.jpg")
img_data = np.zeros_like(img_data)

def to_full_width(s):
    # 创建一个转换表
    half_width = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
    full_width = "０１２３４５６７８９ＡＢＣＤＥＦＧＨＩＪＫＬＭＮＯＰＱＲＳＴＵＶＷＸＹＺａｂｃｄｅｆｇｈｉｊｋｌｍｎｏｐｑｒｓｔｕｖｗｘｙｚ！＂＃＄％＆’（）＊＋，－．／：；＜＝＞？＠［＼］＾＿‘｛｜｝～"

    # 创建映射表
    translate_table = str.maketrans(half_width, full_width)
    
    # 转换字符串
    return s.translate(translate_table)

# 对纸张进行网格化分区, 通过找到最小的行/列, 将其作为网格化的最小单位
x_area_min = np.inf
y_area_min = np.inf
area_lt_point = [np.inf,np.inf]
area_rb_point = [0,0]
for [points, chars] in ocr_result:
    points = np.array(points)
    min_x = min(points[:,0])
    max_x = max(points[:,0])
    min_y = min(points[:,1])
    max_y = max(points[:,1])

    if min_x<area_lt_point[0]:
        area_lt_point[0] = min_x
    if min_y<area_lt_point[1]:
        area_lt_point[1] = min_y

    if max_x>area_rb_point[0]:
        area_rb_point[0] = max_x
    if max_y>area_rb_point[1]:
        area_rb_point[1] = max_y

    if (max_x-min_x)/len(chars[0]) < x_area_min: # 最小x跨度
        x_area_min = (max_x-min_x)/len(chars[0])
    
    if (max_y-min_y) < y_area_min: # 最小y跨度
        y_area_min = max_y-min_y
    
# 行数量与列数量
row_num = math.ceil((area_rb_point[1] - area_lt_point[1])/y_area_min)+1
col_num = math.ceil((area_rb_point[0] - area_lt_point[0])/x_area_min)+1

# list版本的文字图---
text_mat = []
for _ in range(row_num):
    row_list = []
    for _ in range(col_num):
        row_list.append("\u3000")
    text_mat.append(row_list)

for [points, chars] in ocr_result:
    points = np.array(points)
    start_x = min(points[:,0]) - area_lt_point[0]
    start_y = (max(points[:,1])+min(points[:,1]))/2 - area_lt_point[1]
    for idx, char_ in enumerate(chars[0]):
        # if char_ == "历":
        #     print("")
        y = start_y
        x = start_x + idx*x_area_min
        y_idx = math.floor(y/y_area_min)
        x_idx = math.floor(x/x_area_min)
        text_mat[y_idx][x_idx] = to_full_width(char_)
with open("temp.txt","w") as f:
    for str_line in text_mat:
        str_line = "".join(str_line)
        f.write(str_line+"\n")
        # f.flush()
